import cv2
import numpy as np
import matplotlib.pyplot as plt

# 1. 读取图像（灰度模式）
img = cv2.imread('4.bmp', 0)  # 替换为你的图像路径

# 2. 图像预处理：转换为浮点型 + 扩展图像尺寸（DFT要求尺寸为2的幂，避免频谱失真）
rows, cols = img.shape
m = cv2.getOptimalDFTSize(rows)
n = cv2.getOptimalDFTSize(cols)
img_optim = cv2.copyMakeBorder(img, 0, m - rows, 0, n - cols, cv2.BORDER_CONSTANT, value=0)
img_float = np.float32(img_optim)

# 3. 执行傅里叶变换（返回实部和虚部组成的复数矩阵）
dft = cv2.dft(img_float, flags=cv2.DFT_COMPLEX_OUTPUT)

# 4. 频谱中心化（将低频移到图像中心，便于观察）
dft_shift = np.fft.fftshift(dft)

# 5. 计算幅值谱（可视化用）
magnitude = cv2.magnitude(dft_shift[:, :, 0], dft_shift[:, :, 1])
magnitude_log = 20 * np.log(1 + magnitude)  # 对数变换扩展动态范围

# 6. 计算相位谱
phase = cv2.phase(dft_shift[:, :, 0], dft_shift[:, :, 1])

# 7. 逆傅里叶变换（重构图像）
dft_shift_back = np.fft.ifftshift(dft_shift)  # 先逆中心化
img_back = cv2.idft(dft_shift_back)
img_back_magnitude = cv2.magnitude(img_back[:, :, 0], img_back[:, :, 1])
img_back_normalized = cv2.normalize(img_back_magnitude, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

# 8. 结果可视化
plt.figure(figsize=(12, 8))

# 原图
plt.subplot(231)
plt.imshow(img, cmap='gray')
plt.title('Original Image')
plt.axis('off')

# 幅值谱（未对数变换）
plt.subplot(232)
plt.imshow(magnitude, cmap='gray')
plt.title('Magnitude Spectrum')
plt.axis('off')

# 幅值谱（对数变换后）
plt.subplot(233)
plt.imshow(magnitude_log, cmap='gray')
plt.title('Log Magnitude Spectrum')
plt.axis('off')

# 相位谱
plt.subplot(234)
plt.imshow(phase, cmap='gray')
plt.title('Phase Spectrum')
plt.axis('off')

# 重构图像
plt.subplot(235)
plt.imshow(img_back_normalized, cmap='gray')
plt.title('Reconstructed Image')
plt.axis('off')

plt.tight_layout()
plt.show()